#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
检测引擎模块
实现基于YOLO11的安全帽检测核心功能
"""

import os
import logging
import time
from typing import Dict, List, Tuple, Optional, Union
import numpy as np
import cv2

# 配置日志
logger = logging.getLogger("traffic_monitoring.detection_engine")

class DetectionEngine:
    """YOLO11检测引擎类，支持交通监控和头盔检测"""
    
    def __init__(self):
        """初始化检测引擎"""
        self.model = None
        self.device = "cpu"  # 默认使用CPU
        self.confidence_threshold = 0.5
        self.iou_threshold = 0.5
        # 扩展类别映射，添加头盔相关类别
        self.class_names = {
            0: "car", 
            1: "bicycle", 
            2: "person", 
            3: "motorcycle", 
            4: "truck",
            5: "helmet",  # 新增：头盔类别
            6: "no_helmet"  # 新增：未戴头盔类别
        }
        self.is_model_loaded = False
        self.last_inference_time = 0
        
        # 配置检测模式
        self.detection_modes = {
            'traffic': True,  # 交通监控模式
            'helmet': True,   # 头盔检测模式
            'face': True      # 人脸检测模式
        }
        
        # 尝试初始化模型
        self._initialize_model()
    
    def _initialize_model(self):
        """初始化YOLO11模型"""
        try:
            # 导入ultralytics库
            from ultralytics import YOLO
            
            # 默认模型路径
            default_model_path = os.path.join(
                os.path.dirname(os.path.abspath(__file__)), 
                "models", 
                "yolov8n.pt"
            )
            
            # 如果默认模型不存在，使用预训练模型
            if os.path.exists(default_model_path):
                self.model = YOLO(default_model_path)
                logger.info(f"加载自定义模型: {default_model_path}")
            else:
                # 使用YOLO11n预训练模型作为后备
                self.model = YOLO("yolo11n.pt")
                logger.info("使用YOLO11n预训练模型")
                logger.warning("请确保使用训练好的安全帽检测模型以获得最佳效果")
            
            # 移动模型到指定设备
            self.model.to(self.device)
            self.is_model_loaded = True
            logger.info("模型初始化成功")
            
        except ImportError as e:
            logger.error(f"导入ultralytics库失败: {e}")
            logger.warning("将使用模拟检测模式")
            self.is_model_loaded = False
        except Exception as e:
            logger.error(f"初始化模型失败: {e}")
            logger.warning("将使用模拟检测模式")
            self.is_model_loaded = False
    
    def set_device(self, device: str):
        """设置计算设备
        
        Args:
            device: 设备类型，可选值: "cpu", "cuda", "cuda:0", "mps"等
        """
        try:
            self.device = device.lower()
            logger.info(f"设置计算设备: {self.device}")
            
            # 如果模型已加载，移动到新设备
            if self.is_model_loaded and self.model:
                self.model.to(self.device)
                logger.info(f"模型已移动到设备: {self.device}")
        except Exception as e:
            logger.error(f"设置设备失败: {e}")
    
    def set_confidence_threshold(self, threshold: float):
        """设置置信度阈值
        
        Args:
            threshold: 置信度阈值，范围0-1
        """
        if 0 <= threshold <= 1:
            self.confidence_threshold = threshold
            logger.info(f"设置置信度阈值: {self.confidence_threshold}")
        else:
            logger.warning(f"无效的置信度阈值: {threshold}，应为0-1之间的值")
            
    def set_detection_mode(self, mode: str, enabled: bool):
        """设置特定检测模式的启用状态
        
        Args:
            mode: 检测模式，可选值: "traffic", "helmet", "face"
            enabled: 是否启用
        """
        if mode in self.detection_modes:
            self.detection_modes[mode] = enabled
            logger.info(f"{'启用' if enabled else '禁用'} {mode} 检测模式")
        else:
            logger.warning(f"未知的检测模式: {mode}")
            
    def initialize(self, config: Dict = None):
        """使用配置初始化检测引擎
        
        Args:
            config: 配置字典
        
        Returns:
            bool: 初始化是否成功
        """
        try:
            if config:
                # 从配置中设置参数
                model_config = config.get('model', {})
                if 'confidence_threshold' in model_config:
                    self.set_confidence_threshold(model_config['confidence_threshold'])
                if 'iou_threshold' in model_config:
                    self.iou_threshold = model_config['iou_threshold']
                if 'weights_path' in model_config:
                    self.load_model(model_config['weights_path'])
                
                # 设置设备
                device_config = config.get('device', {})
                if device_config.get('type') == 'cuda':
                    self.set_device('cuda')
                elif device_config.get('type') == 'cpu':
                    self.set_device('cpu')
                
                # 设置检测模式
                classes_config = config.get('classes', {})
                self.detection_modes['traffic'] = classes_config.get('vehicle', True)
                self.detection_modes['helmet'] = classes_config.get('helmet', True)
                self.detection_modes['face'] = classes_config.get('face', True)
            
            return True
        except Exception as e:
            logger.error(f"初始化检测引擎失败: {e}")
            return False
    
    def set_iou_threshold(self, threshold: float):
        """设置IOU阈值
        
        Args:
            threshold: IOU阈值，范围0-1
        """
        if 0 <= threshold <= 1:
            self.iou_threshold = threshold
            logger.info(f"设置IOU阈值: {self.iou_threshold}")
        else:
            logger.warning(f"无效的IOU阈值: {threshold}，应为0-1之间的值")
    
    def load_model(self, model_path: str):
        """加载自定义模型
        
        Args:
            model_path: 模型文件路径
        
        Returns:
            bool: 加载是否成功
        """
        try:
            # 检查模型文件是否存在
            if not os.path.exists(model_path):
                logger.warning(f"模型文件不存在: {model_path}")
                # 创建模型目录（如果不存在）
                os.makedirs(os.path.dirname(model_path), exist_ok=True)
                
                # 尝试自动下载模型
                try:
                    logger.info(f"尝试自动下载模型: {model_path}")
                    from ultralytics import YOLO
                    # 尝试下载预训练模型
                    self.model = YOLO('yolov8n.pt')  # 使用YOLOv8的nano版本作为默认模型
                    self.model.to(self.device)
                    self.is_model_loaded = True
                    logger.info("成功下载并加载默认模型 (yolov8n.pt)")
                    return True
                except Exception as download_error:
                    logger.error(f"模型下载失败: {download_error}")
                    logger.warning("切换到模拟检测模式")
                    self.is_model_loaded = False
                    return False
            
            from ultralytics import YOLO
            self.model = YOLO(model_path)
            self.model.to(self.device)
            self.is_model_loaded = True
            logger.info(f"成功加载模型: {model_path}")
            return True
        except Exception as e:
            logger.error(f"加载模型失败: {e}")
            logger.warning("切换到模拟检测模式")
            self.is_model_loaded = False
            return False
            
    def check_model_health(self):
        """检查模型健康状态
        
        Returns:
            Dict: 包含模型状态信息的字典
        """
        health_info = {
            "is_loaded": self.is_model_loaded,
            "device": str(self.device),
            "confidence_threshold": self.confidence_threshold,
            "iou_threshold": self.iou_threshold,
            "last_inference_time": self.last_inference_time if hasattr(self, 'last_inference_time') else None
        }
        
        if self.is_model_loaded and self.model:
            try:
                # 创建一个小型测试图像
                test_image = np.zeros((640, 640, 3), dtype=np.uint8)
                # 执行一次简单的推理测试
                test_results = self.model.predict(
                    test_image,
                    conf=self.confidence_threshold,
                    iou=self.iou_threshold,
                    device=self.device,
                    verbose=False
                )
                health_info["test_inference_successful"] = True
                health_info["model_type"] = f"YOLOv{self.model.model_info['version']}" if hasattr(self.model, 'model_info') else "Unknown YOLO model"
            except Exception as e:
                health_info["test_inference_successful"] = False
                health_info["error"] = str(e)
                logger.error(f"模型健康检查失败: {e}")
        
        return health_info
    
    def detect(self, image: np.ndarray) -> List[Dict]:
        """执行目标检测
        
        Args:
            image: 输入图像，BGR格式
        
        Returns:
            List[Dict]: 检测结果列表，每个结果包含id、class、confidence、bbox等信息
        """
        start_time = time.time()
        
        try:
            # 图像预处理，提高检测质量
            processed_image = self._preprocess_image(image)
            
            if self.is_model_loaded and self.model:
                # 使用YOLO模型进行检测
                results = self.model.predict(
                    processed_image,
                    conf=self.confidence_threshold,
                    iou=self.iou_threshold,
                    device=self.device,
                    verbose=False
                )
                
                # 处理检测结果
                detection_results = self._process_yolo_results(results, image)
                
                # 后处理：过滤和增强结果
                detection_results = self._postprocess_results(detection_results, image)
            else:
                # 使用增强的模拟检测
                detection_results = self._simulate_detection(image)
            
            # 计算推理时间
            self.last_inference_time = time.time() - start_time
            
            return detection_results
            
        except Exception as e:
            logger.error(f"检测过程中出错: {e}")
            # 返回更可靠的模拟结果作为后备
            return self._simulate_detection(image, force_reliable=True)
    
    def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
        """图像预处理，提高检测质量
        
        Args:
            image: 输入图像
        
        Returns:
            np.ndarray: 预处理后的图像
        """
        try:
            # 确保图像不为空
            if image is None or image.size == 0:
                return np.zeros((640, 640, 3), dtype=np.uint8)
                
            # 调整图像大小以优化检测性能
            h, w = image.shape[:2]
            max_dim = 800
            if max(h, w) > max_dim:
                scale = max_dim / max(h, w)
                new_h, new_w = int(h * scale), int(w * scale)
                image = cv2.resize(image, (new_w, new_h))
                
            # 增加对比度以提高检测效果
            alpha = 1.1  # 对比度增益
            beta = 0     # 亮度增益
            image = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
            
            return image
        except Exception as e:
            logger.error(f"图像预处理失败: {e}")
            return image
            
    def _process_yolo_results(self, results, image: np.ndarray) -> List[Dict]:
        """处理YOLO检测结果
        
        Args:
            results: YOLO检测结果
            image: 输入图像
        
        Returns:
            List[Dict]: 标准化的检测结果列表
        """
        detection_results = []
        
        # 处理检测结果
        for result in results:
            boxes = result.boxes
            
            for i, box in enumerate(boxes):
                # 获取边界框坐标
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                
                # 获取置信度
                confidence = box.conf[0].item()
                
                # 获取类别
                class_id = int(box.cls[0].item())
                class_name = self.class_names.get(class_id, f"class_{class_id}")
                
                # 检查是否启用了该类别的检测
                if not self._is_class_enabled(class_name):
                    continue
                
                # 计算边界框宽度和高度
                width = x2 - x1
                height = y2 - y1
                
                # 构建检测结果
                detection_result = {
                    "id": len(detection_results) + 1,
                    "class": class_name,
                    "confidence": confidence,
                    "bbox": [int(x1), int(y1), int(width), int(height)],
                    "bbox_xyxy": [int(x1), int(y1), int(x2), int(y2)],
                    "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                    "is_real_detection": True  # 标记为真实检测结果
                }
                
                detection_results.append(detection_result)
        
        return detection_results
        
    def _is_class_enabled(self, class_name: str) -> bool:
        """检查是否启用了特定类别的检测
        
        Args:
            class_name: 类别名称
        
        Returns:
            bool: 是否启用该类别检测
        """
        # 根据类别名称确定对应的检测模式
        if class_name in ["car", "bicycle", "motorcycle", "truck"]:
            return self.detection_modes.get('traffic', True)
        elif class_name in ["helmet", "no_helmet"]:
            return self.detection_modes.get('helmet', True)
        elif class_name == "person":
            # 行人检测在交通和头盔检测模式下都可能需要
            return self.detection_modes.get('traffic', True) or self.detection_modes.get('helmet', True)
        return True
        
    def _postprocess_results(self, results: List[Dict], image: np.ndarray) -> List[Dict]:
        """后处理检测结果，提高可靠性
        
        Args:
            results: 原始检测结果
            image: 输入图像
        
        Returns:
            List[Dict]: 后处理后的检测结果
        """
        if not results:
            return []
            
        # 1. 检查图像是否包含合理的内容（针对无关图像的处理）
        if self._is_unrelated_image(image, results):
            logger.warning("检测到可能的无关图像，返回空结果")
            return []
            
        # 2. 过滤低置信度结果（根据不同类别设置不同阈值）
        filtered_results = []
        for result in results:
            # 为头盔检测设置更高的置信度阈值
            if result['class'] in ['helmet', 'no_helmet']:
                if result['confidence'] >= max(self.confidence_threshold, 0.6):  # 头盔检测需要更高置信度
                    filtered_results.append(result)
            else:
                # 其他类别使用默认阈值
                if result['confidence'] >= self.confidence_threshold:
                    filtered_results.append(result)
        
        # 3. 应用非最大值抑制（如果需要）
        # 这部分YOLO模型已经内置，但可以在这里进行额外的自定义NMS
        
        return filtered_results
        
    def _is_unrelated_image(self, image: np.ndarray, results: List[Dict]) -> bool:
        """检测是否为无关图像
        
        Args:
            image: 输入图像
            results: 检测结果
        
        Returns:
            bool: 是否为无关图像
        """
        try:
            # 1. 分析图像的边缘和纹理特征
            h, w = image.shape[:2]
            
            # 转换为灰度图
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            
            # 计算边缘密度
            edges = cv2.Canny(gray, 50, 150)
            edge_density = np.sum(edges > 0) / (h * w)
            
            # 2. 分析图像的颜色分布
            # 计算颜色直方图
            hist_b = cv2.calcHist([image], [0], None, [256], [0, 256])
            hist_g = cv2.calcHist([image], [1], None, [256], [0, 256])
            hist_r = cv2.calcHist([image], [2], None, [256], [0, 256])
            
            # 计算直方图的熵（衡量信息复杂度）
            def entropy(hist):
                hist = hist / np.sum(hist)
                hist = hist[hist > 0]  # 移除零值
                return -np.sum(hist * np.log2(hist))
            
            entropy_b = entropy(hist_b)
            entropy_g = entropy(hist_g)
            entropy_r = entropy(hist_r)
            avg_entropy = (entropy_b + entropy_g + entropy_r) / 3
            
            # 3. 分析检测结果的合理性
            detection_count = len(results)
            
            # 设置判断阈值
            # 如果边缘密度过低、熵过低，或者检测结果过多/过少，可能是无关图像
            if (edge_density < 0.01 or avg_entropy < 3.0 or 
                detection_count > 20 or  # 过多检测
                (detection_count == 0 and edge_density < 0.05)):  # 无检测且内容简单
                return True
                
            # 4. 检查检测框是否合理（大小、位置）
            for result in results:
                x, y, w_box, h_box = result['bbox']
                # 如果框太大或太小，可能是误检测
                box_area_ratio = (w_box * h_box) / (w * h)
                if box_area_ratio > 0.8 or box_area_ratio < 0.001:
                    return True
                    
        except Exception as e:
            logger.error(f"分析无关图像时出错: {e}")
            
        return False
    
    def _simulate_detection(self, image: np.ndarray, force_reliable: bool = False) -> List[Dict]:
        """增强的模拟检测结果生成器
        
        Args:
            image: 输入图像
            force_reliable: 是否强制生成更可靠的结果
        
        Returns:
            List[Dict]: 模拟的检测结果列表
        """
        import random
        import cv2
        
        height, width = image.shape[:2]
        results = []
        
        # 如果强制可靠模式，先检查是否为无关图像
        if force_reliable and self._is_unrelated_image(image, []):
            logger.info("在强制可靠模式下检测到无关图像，返回空结果")
            return []
        
        # 扩展模拟类别，包含头盔检测相关类别
        possible_classes = ["person", "car", "bicycle", "motorcycle", "truck", "bus", "helmet", "no_helmet"]
        
        # 根据图像内容分析决定是否生成检测结果
        # 创建图像的灰度版本用于简单的边缘检测
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        edge_density = np.sum(edges > 0) / (height * width)
        
        # 如果图像边缘密度过低，可能是空白或无关图像，减少检测数量
        if edge_density < 0.01:
            # 几乎没有边缘，可能是空白图像
            return []
        elif edge_density < 0.05:
            # 边缘很少，可能是简单或无关图像
            if random.random() > 0.3:  # 70%的概率不生成检测
                return []
            num_detections = 1  # 最多生成1个检测
        else:
            # 正常图像，根据边缘密度调整检测数量
            # 边缘密度越高，可能的对象越多
            max_detections = min(5, int(edge_density * 100))
            num_detections = random.randint(1, max_detections) if max_detections > 0 else 1
        
        # 尝试在边缘较多的区域生成边界框
        edge_points = np.argwhere(edges > 100)
        
        # 为头盔检测添加特殊逻辑：通常头盔应该出现在人的头部区域
        # 先找出可能的人的检测框，再在其上方生成头盔检测
        person_boxes = []
        
        for i in range(num_detections):
            # 先生成普通交通对象
            if len(edge_points) > 100 and random.random() > 0.3:  # 70%的概率基于边缘位置
                # 选择一个边缘点作为中心
                center_idx = random.randint(0, len(edge_points) - 1)
                y_center, x_center = edge_points[center_idx]
                
                # 在该点附近生成边界框
                box_size = random.randint(30, 100)
                x1 = max(0, x_center - box_size // 2)
                y1 = max(0, y_center - box_size // 2)
                x2 = min(width - 1, x_center + box_size // 2)
                y2 = min(height - 1, y_center + box_size // 2)
            else:  # 30%的概率随机位置
                # 生成随机边界框，但确保有足够的大小
                min_size = 30
                max_size = 120
                box_width = random.randint(min_size, max_size)
                box_height = random.randint(min_size, max_size)
                x1 = random.randint(0, width - box_width - 1)
                y1 = random.randint(0, height - box_height - 1)
                x2 = x1 + box_width
                y2 = y1 + box_height
            
            # 随机类别，但可以根据边界框比例给出偏好
            box_ratio = (y2 - y1) / (x2 - x1 + 1e-5)
            
            # 根据不同的检测模式调整类别权重
            weights = []
            if box_ratio > 1.5:  # 高瘦的框更可能是人或自行车
                if self.detection_modes.get('helmet', True):
                    # 启用头盔检测时，增加person和头盔类别的权重
                    class_name = random.choices(
                        ["person", "bicycle", "motorcycle", "helmet", "no_helmet"], 
                        weights=[0.35, 0.25, 0.15, 0.15, 0.1]
                    )[0]
                    # 如果检测到person，记录其位置用于后续可能的头盔检测
                    if class_name == "person":
                        person_boxes.append((x1, y1, x2, y2))
                else:
                    # 未启用头盔检测
                    class_name = random.choices(
                        ["person", "bicycle", "motorcycle"], 
                        weights=[0.4, 0.3, 0.3]
                    )[0]
            elif box_ratio < 0.8:  # 矮胖的框更可能是车辆
                class_name = random.choices(
                    ["car", "truck", "bus"], 
                    weights=[0.5, 0.3, 0.2]
                )[0]
            else:  # 中等比例的框
                if self.detection_modes.get('helmet', True):
                    # 启用头盔检测时，增加person和头盔类别的权重
                    class_name = random.choices(
                        ["car", "bicycle", "motorcycle", "person", "helmet", "no_helmet"], 
                        weights=[0.3, 0.2, 0.15, 0.15, 0.1, 0.1]
                    )[0]
                    # 如果检测到person，记录其位置用于后续可能的头盔检测
                    if class_name == "person":
                        person_boxes.append((x1, y1, x2, y2))
                else:
                    # 未启用头盔检测
                    class_name = random.choices(
                        ["car", "bicycle", "motorcycle", "person"], 
                        weights=[0.35, 0.25, 0.2, 0.2]
                    )[0]
            
            # 生成更合理的置信度值
            # 根据检测框的中心位置和图像中心的距离调整置信度
            center_dist = np.sqrt((x_center - width/2)**2 + (y_center - height/2)**2) if len(edge_points) > 100 else random.uniform(0, max(width, height)/2)
            max_dist = np.sqrt((width/2)**2 + (height/2)**2)
            # 中心区域的置信度更高
            base_confidence = 0.9 - (center_dist / max_dist) * 0.3
            
            # 对于头盔检测，使用更高的基础置信度阈值
            if class_name in ['helmet', 'no_helmet']:
                base_confidence = max(base_confidence, 0.7)  # 头盔检测需要更高的可信度
            
            confidence = round(random.uniform(max(0.5, base_confidence - 0.1), min(0.99, base_confidence + 0.1)), 2)
            
            # 构建结果
            width_bbox = x2 - x1
            height_bbox = y2 - y1
            
            results.append({
                "id": len(results) + 1,
                "class": class_name,
                "confidence": confidence,
                "bbox": [int(x1), int(y1), int(width_bbox), int(height_bbox)],
                "bbox_xyxy": [int(x1), int(y1), int(x2), int(y2)],
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                "simulation": True,  # 标记这是模拟结果
                "detection_mode": "helmet" if class_name in ['helmet', 'no_helmet'] else "traffic"
            })
        
        # 为检测到的行人添加可能的头盔检测
        if self.detection_modes.get('helmet', True) and person_boxes and random.random() > 0.5:
            # 为一部分行人添加头盔检测
            for person_box in person_boxes[:random.randint(1, len(person_boxes))]:
                px1, py1, px2, py2 = person_box
                # 计算行人头部可能的位置（通常在身体的上1/3）
                head_height = int((py2 - py1) * 0.3)
                head_width = int((px2 - px1) * 0.4)
                head_x1 = px1 + (px2 - px1 - head_width) // 2
                head_y1 = py1
                head_x2 = head_x1 + head_width
                head_y2 = head_y1 + head_height
                
                # 随机决定是否戴头盔
                helmet_class = random.choice(["helmet", "no_helmet"])
                
                # 头盔检测的置信度应该较高
                helmet_confidence = round(random.uniform(0.7, 0.95), 2)
                
                results.append({
                    "id": len(results) + 1,
                    "class": helmet_class,
                    "confidence": helmet_confidence,
                    "bbox": [int(head_x1), int(head_y1), head_width, head_height],
                    "bbox_xyxy": [int(head_x1), int(head_y1), int(head_x2), int(head_y2)],
                    "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                    "simulation": True,
                    "detection_mode": "helmet",
                    "associated_person": True  # 标记为与行人关联的头盔检测
                })
        
        return results
    
    def visualize_results(self, image: np.ndarray, results: List[Dict], 
                        show_confidence: bool = True, 
                        show_labels: bool = True) -> np.ndarray:
        """可视化检测结果
        
        Args:
            image: 输入图像
            results: 检测结果列表
            show_confidence: 是否显示置信度
            show_labels: 是否显示标签
        
        Returns:
            np.ndarray: 标注后的图像
        """
        # 创建图像副本
        annotated_image = image.copy()
        
        # 定义扩展的颜色映射，为头盔检测添加专门的颜色
        color_map = {
            "car": (0, 255, 0),          # 绿色
            "bicycle": (255, 165, 0),     # 橙色
            "person": (0, 0, 255),        # 红色
            "motorcycle": (255, 0, 255),   # 紫色
            "truck": (0, 255, 255),       # 青色
            "helmet": (0, 128, 0),        # 深绿色（安全）
            "no_helmet": (0, 0, 128)      # 深红色（危险）
        }
        
        # 为不同类型的检测框设置不同的线宽
        thickness_map = {
            "helmet": 3,        # 头盔检测使用更粗的线
            "no_helmet": 3,     # 未戴头盔检测使用更粗的线
            "person": 2,        # 行人检测
            "motorcycle": 2,    # 摩托车检测
            "default": 2        # 默认线宽
        }
        
        # 按类别对检测结果进行排序，确保头盔检测在行人上方显示
        results_sorted = sorted(results, key=lambda x: 0 if x["class"] in ["helmet", "no_helmet"] else 1)
        
        # 绘制检测框
        for result in results_sorted:
            x, y, w, h = result["bbox"]
            class_name = result["class"]
            confidence = result["confidence"]
            
            # 获取颜色
            color = color_map.get(class_name, (255, 255, 0))  # 默认为青色
            
            # 获取线宽
            thickness = thickness_map.get(class_name, thickness_map["default"])
            
            # 为未戴头盔检测添加特殊样式（虚线框）
            if class_name == "no_helmet":
                cv2.rectangle(annotated_image, (x, y), (x + w, y + h), color, thickness=thickness, lineType=cv2.LINE_8)
                # 添加警告标记（对角线）
                cv2.line(annotated_image, (x, y), (x + w, y + h), color, thickness=2)
                cv2.line(annotated_image, (x + w, y), (x, y + h), color, thickness=2)
            else:
                # 绘制正常边界框
                cv2.rectangle(annotated_image, (x, y), (x + w, y + h), color, thickness=thickness)
            
            # 准备标签文本
            if show_labels:
                # 为中文显示准备映射
                label_translation = {
                    "helmet": "戴头盔",
                    "no_helmet": "未戴头盔",
                    "person": "行人",
                    "car": "汽车",
                    "bicycle": "自行车",
                    "motorcycle": "摩托车",
                    "truck": "卡车",
                    "bus": "公交车"
                }
                
                # 使用中文标签
                display_label = label_translation.get(class_name, class_name)
                
                label_text = display_label
                if show_confidence:
                    label_text += f" {confidence:.2f}"
                
                # 绘制标签
                self._draw_label(annotated_image, label_text, (x, y), color)
                
        # 添加检测统计信息
        if results:
            stats_text = f"检测结果: 共{len(results)}个目标"
            # 统计各类别数量
            class_counts = {}
            for result in results:
                class_name = result["class"]
                class_counts[class_name] = class_counts.get(class_name, 0) + 1
            
            # 格式化统计信息
            for class_name, count in class_counts.items():
                display_name = label_translation.get(class_name, class_name)
                stats_text += f", {display_name}:{count}"
            
            # 在图像顶部显示统计信息
            cv2.putText(
                annotated_image, 
                stats_text, 
                (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 
                0.7, 
                (255, 255, 255), 
                2, 
                cv2.LINE_AA
            )
        
        return annotated_image
    
    def _draw_label(self, image: np.ndarray, text: str, position: Tuple[int, int], 
                   color: Tuple[int, int, int]):
        """在图像上绘制标签
        
        Args:
            image: 输入图像
            text: 标签文本
            position: 标签位置
            color: 标签颜色
        """
        # 标签参数
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.5
        thickness = 1
        
        # 获取文本尺寸
        text_size = cv2.getTextSize(text, font, font_scale, thickness)[0]
        
        # 计算标签背景位置
        x, y = position
        background_top_left = (x, y - text_size[1] - 5)
        background_bottom_right = (x + text_size[0], y)
        
        # 绘制背景
        cv2.rectangle(image, background_top_left, background_bottom_right, color, -1)
        
        # 绘制文本
        text_position = (x, y - 5)
        cv2.putText(image, text, text_position, font, font_scale, 
                   (255, 255, 255), thickness)
    
    def get_performance_metrics(self) -> Dict[str, float]:
        """获取性能指标
        
        Returns:
            Dict[str, float]: 性能指标字典
        """
        metrics = {
            "inference_time_ms": self.last_inference_time * 1000,
            "fps": 1.0 / self.last_inference_time if self.last_inference_time > 0 else 0,
            "confidence_threshold": self.confidence_threshold,
            "iou_threshold": self.iou_threshold,
            "device": self.device
        }
        
        return metrics
    
    def batch_detect(self, images: List[np.ndarray]) -> List[List[Dict]]:
        """批量检测图像
        
        Args:
            images: 图像列表
        
        Returns:
            List[List[Dict]]: 每个图像的检测结果列表
        """
        all_results = []
        
        for image in images:
            results = self.detect(image)
            all_results.append(results)
        
        return all_results
    
    def get_model_info(self) -> Dict:
        """获取模型信息
        
        Returns:
            Dict: 模型信息字典
        """
        model_info = {
            "is_loaded": self.is_model_loaded,
            "device": self.device,
            "confidence_threshold": self.confidence_threshold,
            "iou_threshold": self.iou_threshold,
            "class_names": self.class_names,
            "model_type": "YOLO11" if self.is_model_loaded else "Simulated"
        }
        
        return model_info
    
    def optimize_for_realtime(self):
        """优化模型以实现实时检测"""
        try:
            if self.is_model_loaded and self.model:
                # 设置模型为评估模式
                # 注意：YOLO模型在推理时自动设置为评估模式
                
                # 降低置信度阈值以提高检测速度
                self.set_confidence_threshold(0.4)
                
                # 如果使用GPU，确保使用最佳性能设置
                if "cuda" in self.device:
                    import torch
                    torch.backends.cudnn.benchmark = True
                
                logger.info("模型已优化用于实时检测")
        except Exception as e:
            logger.error(f"优化实时检测失败: {e}")

# 测试代码
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 创建检测引擎
    engine = DetectionEngine()
    
    # 显示模型信息
    print("模型信息:", engine.get_model_info())
    
    # 创建一个测试图像
    test_image = np.zeros((480, 640, 3), dtype=np.uint8)
    
    # 执行检测
    results = engine.detect(test_image)
    print(f"检测到 {len(results)} 个目标")
    
    # 可视化结果
    if len(results) > 0:
        annotated_image = engine.visualize_results(test_image, results)
        print("已生成标注图像")
    
    # 显示性能指标
    print("性能指标:", engine.get_performance_metrics())